import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms

from Dataloader import CustomDataset
from model import SwinDeepLab


def train(model, dataloader, criterion, optimizer, device, num_epochs=31):
    model.train()
    for epoch in range(num_epochs):
        print(f'第{epoch + 1}轮训练:')
        for images, masks in dataloader:
            images = images.float().to(device)
            masks = masks.float().to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()

        print(f'Epoch {epoch}, Loss: {loss.item()}')
        if (epoch + 1) % 3 == 0:
            torch.save(model.state_dict(), f'../pt_file/Unet_ConvTranspose2d_better-{epoch}.pt')


if __name__ == '__main__':
    # 数据集与数据加载器
    # 可以根据需求进行数据增强操作
    transform = transforms.Compose([
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    input_dir = '../unet_img/imgs/val'
    mask_dir = '../unet_img/masks/val'
    dataset = CustomDataset(input_dir, mask_dir, transform=transform)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

    # 初始化模型、损失函数和优化器
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SwinDeepLab(1).to(device)
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 开始训练
    train(model, dataloader, criterion, optimizer, device)
